#include "hnurm_detect/detect_node.hpp"

#include <cv_bridge/cv_bridge.h>
#include <sensor_msgs/image_encodings.hpp>

namespace hnurm
{
DetectNode::DetectNode(const rclcpp::NodeOptions &options) : rclcpp::Node("detect_node", options)
{
    init_params();
    // init detector
    detector_ = std::make_shared<Detector>(min_lightness_, light_params_, armor_params_, classifier_params_);

    // subs
    if (debug_with_no_c_board)
    {
        debug_img_sub = this->create_subscription<sensor_msgs::msg::Image>(
            "image", rclcpp::SensorDataQoS(), std::bind(&DetectNode::debug_detect_callback, this, std::placeholders::_1)
        );
    }
    else{
        sub_img_ = std::make_shared<message_filters::Subscriber<sensor_msgs::msg::Image>>(
            this, image_topic_, rmw_qos_profile_sensor_data
        );
        sub_uart_ = std::make_shared<message_filters::Subscriber<hnurm_interfaces::msg::VisionRecvData>>(
            this, recv_topic_, rmw_qos_profile_sensor_data
        );
        sync_ = std::make_shared<message_filters::Synchronizer<ApproximateTimePolicy>>(
            ApproximateTimePolicy(4'096), *sub_img_, *sub_uart_
        );
        sync_->registerCallback(&DetectNode::detect_callback, this);
    }
    // pubs
    pub_res_img_ = image_transport::create_publisher(this, res_img_topic_, rmw_qos_profile_sensor_data);
    pub_num_img_ = image_transport::create_publisher(this, num_img_topic_, rmw_qos_profile_sensor_data);
    pub_bin_img_ = image_transport::create_publisher(this, bin_img_topic_, rmw_qos_profile_sensor_data);
    pub_res_armor_
        = this->create_publisher<hnurm_interfaces::msg::ArmorArray>(res_armor_topic_, rclcpp::SensorDataQoS());

    // debug pubs
    std::string debug_topic = this->declare_parameter("debug_info_topic", "/debug/detect_time");
    pub_detect_time_        = this->create_publisher<std_msgs::msg::Float64>(debug_topic, rclcpp::SensorDataQoS());
}

void DetectNode::init_params()
{
    auto f =
        // topics
        recv_topic_  = this->declare_parameter("vision_recv_topic", "vision_recv_data");
    image_topic_     = this->declare_parameter("img_topic", "image");
    res_img_topic_   = this->declare_parameter("res_img_topic", "res_img");
    num_img_topic_   = this->declare_parameter("num_img_topic", "num_img");
    bin_img_topic_   = this->declare_parameter("bin_img_topic", "bin_img");
    res_armor_topic_ = this->declare_parameter("res_armor_topic", "armor");

    pub_debug_info_ = this->declare_parameter("pub_debug_info.pub_debug_info", false);
    pub_debug_res_img = this->declare_parameter("pub_debug_info.pub_debug_res_img", false);

    pub_debug_num_img = this->declare_parameter("pub_debug_info.pub_debug_num_img", false);

    pub_debug_bin_img = this->declare_parameter("pub_debug_info.pub_debug_bin_img", false);
    width = this->declare_parameter("pub_debug_info.width", 320);
    height = this->declare_parameter("pub_debug_info.height", 320);
    use_ros = this->declare_parameter("pub_debug_info.use_ros", true);
    debug_with_no_c_board = this->declare_parameter("pub_debug_info.debug_with_no_c_board",false);
    debug_self_color = this->declare_parameter("pub_debug_info.debug_self_color","blue");

    // detector params
    min_lightness_ = static_cast<int>(this->declare_parameter("min_lightness", 120));

    light_params_.min_ratio = this->declare_parameter("min_ratio", 0.12);
    light_params_.max_ratio = this->declare_parameter("max_ratio", 0.55);
    light_params_.max_angle = this->declare_parameter("light_max_angle", 40.0);

    armor_params_.min_light_ratio           = this->declare_parameter("min_light_ratio", 0.8);
    armor_params_.min_small_center_distance = this->declare_parameter("min_small_center_distance", 0.8);
    armor_params_.max_small_center_distance = this->declare_parameter("max_small_center_distance", 3.0);
    armor_params_.min_large_center_distance = this->declare_parameter("min_large_center_distance", 3.2);
    armor_params_.max_large_center_distance = this->declare_parameter("max_large_center_distance", 4.5);
    armor_params_.max_angle                 = this->declare_parameter("armor_max_angle", 35.0);

    classifier_params_.model_path     = this->declare_parameter("model_path", "mlp.onnx");
    classifier_params_.label_path     = this->declare_parameter("label_path", "label.txt");
    classifier_params_.threshold      = this->declare_parameter("threshold", 0.7);
    classifier_params_.ignore_classes = this->declare_parameter("ignore_classes", std::vector<std::string>());
    if (!use_ros && pub_debug_info_)
    {   if(pub_debug_res_img) cv::namedWindow("res_img");
        if(pub_debug_num_img) cv::namedWindow("num_img");
        if(pub_debug_bin_img)cv::namedWindow("bin_img");
    }
}

void DetectNode::detect_callback(
    const sensor_msgs::msg::Image::SharedPtr img_msg, const hnurm_interfaces::msg::VisionRecvData::SharedPtr uart_msg
)
{
    static auto last_time = this->now();
    auto        this_time = this->now();
    auto        call_dt   = (this_time - last_time).seconds();
    RCLCPP_INFO(this->get_logger(), "Callback FPS: %f", 1.0 / call_dt);
    last_time = this_time;

    // first get cloud
    // 串口中，0==无，1==红，2==蓝
    // 检测任务中，0==红，1==蓝，2==紫
    auto self_color = uart_msg->self_color.data;
    if(self_color == hnurm_interfaces::msg::SelfColor::RED)
        self_color = hnurm::RED;
    else
        self_color = hnurm::BLUE;

    // get image
    static cv_bridge::CvImageConstPtr cv_ptr;
    try
    {
        cv_ptr = cv_bridge::toCvShare(img_msg, sensor_msgs::image_encodings::BGR8);
    }
    catch(cv_bridge::Exception &e)
    {
        RCLCPP_ERROR(this->get_logger(), "cv_bridge exception: %s", e.what());
        return;
    }

    // do detect
    rclcpp::Time start      = this->now();
    static long  detect_cnt = 0;
    detector_->detect(cv_ptr->image, self_color);
    auto dt = (this->now() - start).seconds();


    // pub res
    std_msgs::msg::Float64 dt_;
    dt_.data = dt * 1000.0;
    pub_detect_time_->publish(dt_);

    // publish res images
    if(pub_debug_info_)
    {   
        if (!use_ros) 
        {
            if (pub_debug_res_img)
            {
            cv::Mat res_img = cv_ptr->image.clone();
            detector_->drawResults(res_img);
            cv::resize(res_img, res_img, cv::Size(width, height)); 
            cv::imshow("res_img",res_img);
            }
            if (pub_debug_num_img)
            {
                cv::Mat num_img = detector_->getAllNumbersImage().clone();
                cv::resize(num_img, num_img, cv::Size(width, height)); 
                cv::imshow("num_img",num_img);
            }
            if (pub_debug_bin_img)
            {   
                cv::Mat bin_img = detector_->binary_img;
                cv::resize(bin_img, bin_img, cv::Size(width, height)); 
                cv::imshow("bin_img",bin_img);
                
            } 
            cv::waitKey(1); // 等待1ms，刷新窗口
        }
        else 
        {
            if (pub_debug_res_img)
            {
                cv::Mat res_img = cv_ptr->image.clone();
                detector_->drawResults(res_img);
                auto res_img_msg
                    = cv_bridge::CvImage(std_msgs::msg::Header(), sensor_msgs::image_encodings::BGR8, res_img).toImageMsg();
                pub_res_img_.publish(*res_img_msg);
            }
            if (pub_debug_num_img)
            {
                // publish num images
                cv::Mat num_img = detector_->getAllNumbersImage().clone();
                auto num_img_msg
                    = cv_bridge::CvImage(std_msgs::msg::Header(), sensor_msgs::image_encodings::MONO8, num_img).toImageMsg();
                pub_num_img_.publish(*num_img_msg);
            }
            if (pub_debug_bin_img)
            {   
                // publish binary images
                cv::Mat bin_img = detector_->binary_img;
                auto bin_img_msg
                    = cv_bridge::CvImage(std_msgs::msg::Header(), sensor_msgs::image_encodings::MONO8, bin_img).toImageMsg();
                pub_bin_img_.publish(*bin_img_msg);
            } 
        }
    }

    // publish armors
    hnurm_interfaces::msg::ArmorArray armor_msg;
    armor_msg.recv_uart       = *uart_msg;
    armor_msg.header.stamp    = now();
    armor_msg.header.frame_id = "armor";

    auto armors = detector_->getArmor();
    for(const auto &armor : armors)
    {
        armor_msg.armors.push_back(toROSMsg(armor));
    }
    pub_res_armor_->publish(armor_msg);
}
void DetectNode::debug_detect_callback(
    const sensor_msgs::msg::Image::SharedPtr img_msg
)
{
    static auto last_time = this->now();
    auto        this_time = this->now();
    auto        call_dt   = (this_time - last_time).seconds();
    RCLCPP_INFO(this->get_logger(), "Callback FPS: %f", 1.0 / call_dt);
    last_time = this_time;

    // first get cloud
    // 串口中，0==无，1==红，2==蓝
    // 检测任务中，0==红，1==蓝，2==紫
    // auto self_color = uart_msg->self_color.data;
    int self_color;
    if(debug_self_color == "red")
        self_color = hnurm::RED;
    else
        self_color = hnurm::BLUE;
    RCLCPP_INFO(this->get_logger(), "Self color: %d,debug_self_color:%s",self_color,debug_self_color.c_str());
    // get image
    static cv_bridge::CvImageConstPtr cv_ptr;
    try
    {
        cv_ptr = cv_bridge::toCvShare(img_msg, sensor_msgs::image_encodings::BGR8);
    }
    catch(cv_bridge::Exception &e)
    {
        RCLCPP_ERROR(this->get_logger(), "cv_bridge exception: %s", e.what());
        return;
    }

    // do detect
    rclcpp::Time start      = this->now();
    static long  detect_cnt = 0;
    detector_->detect(cv_ptr->image, self_color);
    auto dt = (this->now() - start).seconds();


    // pub res
    std_msgs::msg::Float64 dt_;
    dt_.data = dt * 1000.0;
    pub_detect_time_->publish(dt_);

    // publish res images
    if(pub_debug_info_)
    {   
        if (!use_ros) 
        {
            if (pub_debug_res_img)
            {
            cv::Mat res_img = cv_ptr->image.clone();
            detector_->drawResults(res_img);
            cv::resize(res_img, res_img, cv::Size(width, height)); 
            cv::imshow("res_img",res_img);
            }
            if (pub_debug_num_img)
            {
                cv::Mat num_img = detector_->getAllNumbersImage().clone();
                cv::resize(num_img, num_img, cv::Size(width, height)); 
                cv::imshow("num_img",num_img);
            }
            if (pub_debug_bin_img)
            {   
                cv::Mat bin_img = detector_->binary_img;
                cv::resize(bin_img, bin_img, cv::Size(width, height)); 
                cv::imshow("bin_img",bin_img);
                
            } 
            cv::waitKey(1); // 等待1ms，刷新窗口
        }
        else 
        {
            if (pub_debug_res_img)
            {
                cv::Mat res_img = cv_ptr->image.clone();
                detector_->drawResults(res_img);
                auto res_img_msg
                    = cv_bridge::CvImage(std_msgs::msg::Header(), sensor_msgs::image_encodings::BGR8, res_img).toImageMsg();
                pub_res_img_.publish(*res_img_msg);
            }
            if (pub_debug_num_img)
            {
                // publish num images
                cv::Mat num_img = detector_->getAllNumbersImage().clone();
                auto num_img_msg
                    = cv_bridge::CvImage(std_msgs::msg::Header(), sensor_msgs::image_encodings::MONO8, num_img).toImageMsg();
                pub_num_img_.publish(*num_img_msg);
            }
            if (pub_debug_bin_img)
            {   
                // publish binary images
                cv::Mat bin_img = detector_->binary_img;
                auto bin_img_msg
                    = cv_bridge::CvImage(std_msgs::msg::Header(), sensor_msgs::image_encodings::MONO8, bin_img).toImageMsg();
                pub_bin_img_.publish(*bin_img_msg);
            } 
        }
    }

    // publish armors
    hnurm_interfaces::msg::ArmorArray armor_msg;
    // armor_msg.recv_uart       = *uart_msg;
    armor_msg.header.stamp    = now();
    armor_msg.header.frame_id = "armor";

    auto armors = detector_->getArmor();
    for(const auto &armor : armors)
    {
        armor_msg.armors.push_back(toROSMsg(armor));
    }
    pub_res_armor_->publish(armor_msg);
}

hnurm_interfaces::msg::Armor toROSMsg(const Armor &armor)
{
    hnurm_interfaces::msg::Armor msg;
    msg.left_light.color      = armor.left_light.color;
    msg.left_light.top.x      = armor.left_light.top.x;
    msg.left_light.top.y      = armor.left_light.top.y;
    msg.left_light.center.x   = armor.left_light.center.x;
    msg.left_light.center.y   = armor.left_light.center.y;
    msg.left_light.bottom.x   = armor.left_light.bottom.x;
    msg.left_light.bottom.y   = armor.left_light.bottom.y;
    msg.left_light.size.x     = armor.left_light.size.width;
    msg.left_light.size.y     = armor.left_light.size.height;
    msg.left_light.length     = armor.left_light.length;
    msg.left_light.width      = armor.left_light.width;
    msg.left_light.tilt_angle = armor.left_light.tilt_angle;
    msg.left_light.angle      = armor.left_light.angle;

    msg.right_light.color      = armor.right_light.color;
    msg.right_light.top.x      = armor.right_light.top.x;
    msg.right_light.top.y      = armor.right_light.top.y;
    msg.right_light.center.x   = armor.right_light.center.x;
    msg.right_light.center.y   = armor.right_light.center.y;
    msg.right_light.bottom.x   = armor.right_light.bottom.x;
    msg.right_light.bottom.y   = armor.right_light.bottom.y;
    msg.right_light.size.x     = armor.right_light.size.width;
    msg.right_light.size.y     = armor.right_light.size.height;
    msg.right_light.length     = armor.right_light.length;
    msg.right_light.width      = armor.right_light.width;
    msg.right_light.tilt_angle = armor.right_light.tilt_angle;
    msg.right_light.angle      = armor.right_light.angle;

    msg.center.x = armor.center.x;
    msg.center.y = armor.center.y;

    msg.points2d.resize(armor.points2d.size());
    for(int j = 0; j < armor.points2d.size(); ++j)
    {
        msg.points2d[j].x = armor.points2d[j].x;
        msg.points2d[j].y = armor.points2d[j].y;
    }
    msg.number_img = *cv_bridge::CvImage(std_msgs::msg::Header(), sensor_msgs::image_encodings::MONO8, armor.number_img)
                          .toImageMsg();

    msg.number                = armor.number;
    msg.idx                   = armor.idx;
    msg.similarity            = armor.similarity;
    msg.confidence            = armor.confidence;
    msg.classification_result = armor.classification_result;
    msg.armor_type.data       = armor.armor_type;

    return msg;
}

}  // namespace hnurm
